/**
 * 
 */
package p;

import java.io.IOException;
import java.io.StringReader;
import java.io.StringWriter;
import java.math.BigInteger;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.cert.X509Certificate;
import java.util.Arrays;
import java.util.Date;
import java.util.Iterator;
import java.util.Vector;

import org.bouncycastle.asn1.x509.BasicConstraints;
import org.bouncycastle.asn1.x509.ExtendedKeyUsage;
import org.bouncycastle.asn1.x509.KeyPurposeId;
import org.bouncycastle.asn1.x509.X509Extensions;
import org.bouncycastle.asn1.x509.X509Name;
import org.bouncycastle.jce.PKCS10CertificationRequest;
import org.bouncycastle.jce.PrincipalUtil;
import org.bouncycastle.jce.X509KeyUsage;
import org.bouncycastle.jce.X509Principal;
import org.bouncycastle.openssl.PEMReader;
import org.bouncycastle.openssl.PEMWriter;
import org.bouncycastle.x509.X509V3CertificateGenerator;
import org.bouncycastle.x509.extension.AuthorityKeyIdentifierStructure;
import org.bouncycastle.x509.extension.SubjectKeyIdentifierStructure;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.DatastoreServiceFactory;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.FetchOptions.Builder;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.Query;
import com.google.appengine.api.datastore.Query.FilterOperator;
import com.google.appengine.api.datastore.Text;

/**
 * @author Panxiaobo
 * 
 */
public class CertManager {

	private String toPEM(Object o) throws IOException {
		StringWriter sw = new StringWriter();
		PEMWriter p = new PEMWriter(sw);
		p.writeObject(o);
		p.flush();
		p.close();
		sw.flush();
		return sw.toString();
	}

	private Object fromPEM(String s) throws IOException {
		return new PEMReader(new StringReader(s)).readObject();
	}

	private final DatastoreService dss = DatastoreServiceFactory.getDatastoreService();

	private X509Certificate defaultX509Certificate;
	private PrivateKey defaultX509CertificateKey;

	private static KeyPair gen(String alg, int keysize) throws NoSuchAlgorithmException, NoSuchProviderException {
		KeyPairGenerator kpg = KeyPairGenerator.getInstance(alg, "BC");
		kpg.initialize(keysize);
		return kpg.generateKeyPair();
	}

	public X509Certificate getDefaultCA() throws Exception {

		if (defaultX509Certificate != null) {
			return defaultX509Certificate;
		}
		X509Certificate certificate;
		PrivateKey privateKey;
		Iterator<Entity> it = dss.prepare(new Query("CertCA").addFilter("default", FilterOperator.EQUAL, true))
				.asIterable(Builder.withLimit(1)).iterator();
		if (it.hasNext()) {
			Key caKey = (Key) it.next().getProperty("ca");
			Entity e = dss.get(caKey);
			certificate = (X509Certificate) fromPEM(((Text) e.getProperty("encoded")).getValue());
			Object o = fromPEM(((Text) e.getProperty("privateKey")).getValue());
			if (o instanceof KeyPair) {
				privateKey = ((KeyPair) o).getPrivate();
			} else {
				privateKey = (PrivateKey) o;
			}
		} else {
			Key caKey = dss.allocateIds("Cert", 1).getStart();
			String subject = "C=CN, O=Pxb Test, OU=Test Unit,CN=Pxb Test CA";
			String algorithm = "SHA512WITHECDSA";
			Date notBefore = new Date(System.currentTimeMillis());
			Date notAfter = new Date(System.currentTimeMillis() + (1000L * 60 * 60 * 24 * 30 * 12));

			KeyPair kp = gen("ECDSA", 521);
			X509V3CertificateGenerator v3CertGen = new X509V3CertificateGenerator();
			v3CertGen.setSerialNumber(BigInteger.valueOf(caKey.getId()));
			v3CertGen.setIssuerDN(new X509Principal(subject));
			v3CertGen.setNotBefore(notBefore);
			v3CertGen.setNotAfter(notAfter);
			v3CertGen.setSubjectDN(new X509Principal(subject));
			v3CertGen.setPublicKey(kp.getPublic());
			v3CertGen.setSignatureAlgorithm(algorithm);
			v3CertGen.addExtension(X509Extensions.SubjectKeyIdentifier, false,
					new SubjectKeyIdentifierStructure(kp.getPublic()));
			v3CertGen.addExtension(X509Extensions.AuthorityKeyIdentifier, false,
					new AuthorityKeyIdentifierStructure(kp.getPublic()));
			v3CertGen.addExtension(X509Extensions.BasicConstraints, true, new BasicConstraints(0));
			certificate = v3CertGen.generate(kp.getPrivate());

			privateKey = kp.getPrivate();
			String encodedS = toPEM(certificate);
			String privateKeyS = toPEM(privateKey);

			{
				Entity e = new Entity(caKey);
				e.setProperty("subject", subject);
				e.setProperty("notBefore", notBefore);
				e.setProperty("notAfter", notAfter);
				e.setProperty("isCa", true);
				e.setUnindexedProperty("encoded", new Text(encodedS));
				e.setUnindexedProperty("privateKey", new Text(privateKeyS));
				dss.put(e);
			}

			Key toCaKey = dss.allocateIds("CertCA", 1).getStart();
			{
				Entity e = new Entity(toCaKey);
				e.setProperty("default", true);
				e.setUnindexedProperty("ca", caKey);
				dss.put(e);
			}
		}
		defaultX509Certificate = certificate;
		defaultX509CertificateKey = privateKey;
		return certificate;
	}

	private static final Logger log = LoggerFactory.getLogger(CertManager.class);

	@SuppressWarnings("unchecked")
	public X509Certificate sign(PKCS10CertificationRequest csr, boolean server) throws Exception {

		X509Name subject = csr.getCertificationRequestInfo().getSubject();

		log.debug("CSR Subject {}", subject);

		X509Certificate ca = defaultX509Certificate;

		Date notBefore = new Date(System.currentTimeMillis());
		Date notAfter = new Date(System.currentTimeMillis() + (1000L * 60 * 60 * 24 * 30 * 12));

		Key key = dss.allocateIds("Cert", 1).getStart();
		BigInteger serialNumber = BigInteger.valueOf(key.getId());
		log.debug("Cert SerialNumber {}", serialNumber);

		X509V3CertificateGenerator v3CertGen = new X509V3CertificateGenerator();
		v3CertGen.setSerialNumber(serialNumber);
		v3CertGen.setIssuerDN(PrincipalUtil.getSubjectX509Principal(ca));
		v3CertGen.setNotBefore(notBefore);
		v3CertGen.setNotAfter(notAfter);
		v3CertGen.setSubjectDN(subject);
		v3CertGen.setPublicKey(csr.getPublicKey());
		v3CertGen.setSignatureAlgorithm(ca.getSigAlgName());
		v3CertGen.addExtension(X509Extensions.SubjectKeyIdentifier, false,
				new SubjectKeyIdentifierStructure(ca.getPublicKey()));
		v3CertGen.addExtension(X509Extensions.AuthorityKeyIdentifier, false,
				new AuthorityKeyIdentifierStructure(ca.getPublicKey()));
		v3CertGen.addExtension(X509Extensions.BasicConstraints, true, new BasicConstraints(false));
		v3CertGen.addExtension(
				X509Extensions.ExtendedKeyUsage,
				false,
				new ExtendedKeyUsage(new Vector(Arrays.asList(server ? KeyPurposeId.id_kp_serverAuth
						: KeyPurposeId.id_kp_clientAuth))));// ,
		// KeyPurposeId.id_kp_emailProtection
		v3CertGen.addExtension(X509Extensions.KeyUsage, false, new X509KeyUsage(X509KeyUsage.digitalSignature
				| X509KeyUsage.keyAgreement | X509KeyUsage.dataEncipherment | X509KeyUsage.keyAgreement
				| X509KeyUsage.keyEncipherment));

		X509Certificate certificate = v3CertGen.generate(defaultX509CertificateKey);
		String encodedS = toPEM(certificate);

		Entity e = new Entity(key);
		e.setProperty("subject", subject.toString());
		e.setProperty("notBefore", notBefore);
		e.setProperty("notAfter", notAfter);
		e.setProperty("isCa", false);
		e.setProperty("isServer", server);
		e.setProperty("isClient", !server);
		e.setUnindexedProperty("encoded", new Text(encodedS));

		dss.put(e);

		return certificate;
	}

}
